import os
import hashlib
from PyQt5.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QProgressBar, QTextEdit, QFileDialog
from PyQt5.QtCore import Qt, QThread, pyqtSignal

class ScanThread(QThread):
    progress_updated = pyqtSignal(int)
    result_found = pyqtSignal(str)
    scan_complete = pyqtSignal()

    def __init__(self, path):
        super().__init__()
        self.path = path
        self.stop_scan = False
        
        # 模拟恶意软件哈希数据库
        self.malware_hashes = {
            'e58636a9f8d84a7e9a21443c665170a3': '恶意软件样本A',
            'd41d8cd98f00b204e9800998ecf8427e': '恶意软件样本B',  # 空文件哈希
            '5f4dcc3b5aa765d61d8327deb882cf99': '恶意软件样本C'   # 'password'的MD5哈希
        }

    def run(self):
        file_count = 0
        total_files = self.count_files(self.path)

        for root, dirs, files in os.walk(self.path):
            if self.stop_scan:
                break
            
            for file in files:
                if self.stop_scan:
                    break
                
                file_path = os.path.join(root, file)
                file_count += 1
                
                # 更新进度
                progress = int((file_count / total_files) * 100) if total_files > 0 else 0
                self.progress_updated.emit(progress)
                
                # 计算文件哈希
                file_hash = self.calculate_hash(file_path)
                
                # 检查是否为恶意软件
                if file_hash in self.malware_hashes:
                    self.result_found.emit(f'发现恶意软件: {self.malware_hashes[file_hash]}\n路径: {file_path}\n哈希: {file_hash}\n')

        self.scan_complete.emit()

    def count_files(self, path):
        count = 0
        for root, dirs, files in os.walk(path):
            count += len(files)
        return count

    def calculate_hash(self, file_path):
        try:
            hasher = hashlib.md5()
            with open(file_path, 'rb') as f:
                while chunk := f.read(4096):
                    hasher.update(chunk)
            return hasher.hexdigest()
        except Exception as e:
            return f'错误: {str(e)}'

    def stop(self):
        self.stop_scan = True

class MalwareScanner(QWidget):
    def __init__(self):
        super().__init__()
        self.scan_thread = None
        self.init_ui()

    def init_ui(self):
        layout = QVBoxLayout()
        layout.setAlignment(Qt.AlignTop)

        # 选择扫描路径
        path_layout = QHBoxLayout()
        self.path_label = QLabel('扫描路径: 未选择')
        browse_btn = QPushButton('浏览...')
        browse_btn.clicked.connect(self.browse_path)
        
        path_layout.addWidget(self.path_label)
        path_layout.addWidget(browse_btn)
        layout.addLayout(path_layout)

        # 扫描控制按钮
        btn_layout = QHBoxLayout()
        self.start_btn = QPushButton('开始扫描')
        self.start_btn.clicked.connect(self.start_scan)
        
        self.stop_btn = QPushButton('停止扫描')
        self.stop_btn.clicked.connect(self.stop_scan)
        self.stop_btn.setEnabled(False)
        
        btn_layout.addWidget(self.start_btn)
        btn_layout.addWidget(self.stop_btn)
        layout.addLayout(btn_layout)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        layout.addWidget(self.progress_bar)

        # 扫描结果
        result_label = QLabel('扫描结果:')
        layout.addWidget(result_label)
        
        self.result_text = QTextEdit()
        self.result_text.setReadOnly(True)
        layout.addWidget(self.result_text)

        self.setLayout(layout)
        self.selected_path = ''

    def browse_path(self):
        path = QFileDialog.getExistingDirectory(self, '选择扫描目录')
        if path:
            self.selected_path = path
            self.path_label.setText(f'扫描路径: {path}')

    def start_scan(self):
        if not self.selected_path:
            self.result_text.append('请先选择扫描路径')
            return

        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.progress_bar.setValue(0)
        self.result_text.clear()
        self.result_text.append(f'开始扫描: {self.selected_path}')

        self.scan_thread = ScanThread(self.selected_path)
        self.scan_thread.progress_updated.connect(self.update_progress)
        self.scan_thread.result_found.connect(self.add_result)
        self.scan_thread.scan_complete.connect(self.scan_finished)
        self.scan_thread.start()

    def stop_scan(self):
        if self.scan_thread and self.scan_thread.isRunning():
            self.scan_thread.stop()
            self.result_text.append('扫描已停止')

    def update_progress(self, value):
        self.progress_bar.setValue(value)

    def add_result(self, result):
        self.result_text.append(result)

    def scan_finished(self):
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.progress_bar.setValue(100)
        self.result_text.append('扫描完成')